# coding = utf-8

'''
进行聚类
'''

from sklearn.cluster import KMeans
import numpy as np
import joblib


def cluster():
    #读取数据
    items = []
    with open("feature", "r") as file:
        for line in file:
          data = line.strip().split("\t")
          vector = [float(data[0]), float(data[1]), float(data[2]), float(data[3]), float(data[4]), float(data[5])]
          if data[-1] == "train":
              continue
          items.append(vector)

    items = np.array(items)
    kmeans = KMeans(n_clusters=3, random_state=0).fit(items)
    joblib.dump(kmeans, "kmeans_3.model")

    train_result = []
    valid_result = []
    with open("feature", "r") as file:
        for line in file:
            data = line.strip().split("\t")
            vector = [float(data[0]), float(data[1]), float(data[2]), float(data[3]), float(data[4]), float(data[5])]
            vector = np.array(vector)
            vector = vector.reshape((1, len(vector)))
            predict = kmeans.predict(vector)[0]
            if data[-1] == "valid":
                train_result.append(predict)
            else:
                valid_result.append(predict)
    train_result = np.array(train_result)
    valid_result = np.array(valid_result)

    train_split = []
    valid_split = []
    for i in range(np.max(train_result)+1):
        train_split.append((train_result == i).sum())
        valid_split.append((valid_result == i).sum())
    print(train_split)
    print(valid_split)









if __name__ == '__main__':
    cluster()